// xp_loadlibrary.cpp
// ------------------
// Example extended stored procedure DLL for loading libraries (other DLLs)
// out of a database table.
//
// This example makes use of the MemoryModule library located here:
// https://github.com/fancycode/MemoryModule
// The files have been slightly modified to use strict multi-byte, as it had
// some conflicts when introduced into a TCHAR project that could swap over
// to UNICODE.  A copy of the original source files, including the license
// can be found in the archive of this project.

#include "xp_loadlibrary.h"

BOOL APIENTRY DllMain(HMODULE hModule,DWORD  ul_reason_for_call,LPVOID lpReserved)
{
	switch (ul_reason_for_call)
	{
	case DLL_PROCESS_ATTACH:
	case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
	case DLL_PROCESS_DETACH:
		break;
	}
	return TRUE;
}

// Versioning export
__declspec(dllexport) ULONG __GetXpVersion()
{
   return ODS_VERSION;
}

// Xsproc entry point
__declspec(dllexport) SRVRETCODE xp_loadlibrary(SRV_PROC *srvproc)
{
	SRVRETCODE	vRetCode;

	// Call the function and return
	vRetCode = SqlLoadLibrary(srvproc);

	return(vRetCode);
}

// Send a SQL msg (non-error)
SRVRETCODE Sql_SendMsg(SRV_PROC *srvproc,TCHAR* tczMsg)
{
	int	iResult = 0;
	// Both the ASCII and UNICODE calls are assuming a NULL terminated string
#ifdef UNICODE
	// wcslen() is used here because 'SRV_NULLTERM' seems to always result in FAILURE with srv_wsendmsg()
	iResult = srv_wsendmsg(srvproc,(DBTINYINT)0,(DBTINYINT)0,tczMsg,wcslen(tczMsg));//SRV_NULLTERM);
#else
	iResult = srv_sendmsg(srvproc, SRV_MSG_INFO, 0,(DBTINYINT)0,(DBTINYINT)0,NULL,0,0,tczMsg,strlen((CHAR*)tczMsg));
#endif
	if(iResult != SUCCEED)
		return(Sql_SendErr(srvproc,_T("Sql_SendMsg failure.")));
	return(SUCCEED);
}

// Send a SQL msg (error)
SRVRETCODE Sql_SendErr(SRV_PROC *srvproc,TCHAR* tczErr)
{
	int	iResult = 0;
	// Both the ASCII and UNICODE calls are assuming a NULL terminated string
#ifdef UNICODE
	// wcslen() is used here because 'SRV_NULLTERM' seems to always result in FAILURE with srv_wsendmsg()
	iResult = srv_wsendmsg(srvproc,0,SRV_MSG_ERROR,tczErr,wcslen(tczErr));//SRV_NULLTERM);
#else
	iResult = srv_sendmsg(srvproc, SRV_MSG_ERROR, 0,(DBTINYINT)0,(DBTINYINT)0,NULL,0,0,tczErr,strlen((CHAR*)tczErr));
#endif
	if(SUCCEED != iResult)
		return(FAIL);
	return(SUCCEED);
}

// Loads a DLL from a database table
// This is a big ugly function, but it's PoC and
// mostly to show the logic - I warned ya!
SRVRETCODE SqlLoadLibrary(SRV_PROC *srvproc)
{
	int				iLSecure			= 0;
	int				iReturn				= 0;
	int				iDSNLen				= 0;

	TCHAR*			tcUserName			= NULL;
	TCHAR*			tcPassWord			= NULL;
	TCHAR*			tcDSN				= NULL;
	TCHAR*			tcTableName			= NULL;
	TCHAR*			tcLibrary			= NULL;
	TCHAR*			tcFunction			= NULL;
	CHAR*			acFunction			= NULL;
	TCHAR*			tcParameters		= NULL;
	DWORD			dwCharsConverted	= 0;
	BOOL			bImpersonated		= false;

	HENV			henv				= SQL_NULL_HENV;
    HDBC			hdbc				= SQL_NULL_HDBC;
    HSTMT			hstmt				= SQL_NULL_HSTMT;

	SQLSMALLINT		iOutConnStrLen		= 0;

	TCHAR*			tcSqlStmnt			= NULL;

	BYTE*			bpDLLData			= NULL;
	BYTE*			bpDLLTemp			= NULL;
	BYTE*			bpSqlData			= NULL;
	DWORD			dwDLLSize			= 0;
	SQLINTEGER		iBinaryLenOrInd		= 0;
	SQLINTEGER		iNumBytes			= 0;

	HMEMORYMODULE	module				= NULL;

	CHAR**			DllReturnedCharPp	= NULL;

	// Make sure the correct number of parameters was provided
	if(3 > srv_rpcparams(srvproc))
	{
		PrintHelpInfo(srvproc);
		return(XP_ERROR);
	}

	// - Connect back in to SQL -
	
	// Check the authentication type
	srv_pfield(srvproc,SRV_LSECURE,&iLSecure);
	if(6 == iLSecure)
	{
		// Not using integrated auth, get a password for use
		GetPfieldInfo(srvproc,SRV_PWD,&tcPassWord);

		// Get user name for connection
		GetPfieldInfo(srvproc,SRV_USER,&tcUserName);
	}
	else
	{
		// Client is using integrated auth, impersonate
		bImpersonated = srv_impersonate_client(srvproc);
		if(!bImpersonated)
		{
			// Failed to impersonate
			Sql_SendErr(srvproc,_T("Failed to impersonate client."));
			goto CleanUp;
		}
	}

	// Build the connection string
	iDSNLen = _tcslen(_T("Driver={SQL Native Client};Server=127.0.0.1"));
	if(bImpersonated)
		iDSNLen += _tcslen(_T(";Trusted_Connection=yes"));
	else
	{
		iDSNLen += _tcslen(_T(";Uid=;Pwd="));
		iDSNLen += _tcslen(tcUserName);
		iDSNLen += _tcslen(tcPassWord);
	}
	iDSNLen += 1;
	tcDSN = new TCHAR[iDSNLen];
	memset(tcDSN,0,sizeof(TCHAR) * iDSNLen);
	_tcscpy(tcDSN,_T("Driver={SQL Native Client};Server=127.0.0.1"));
	if(bImpersonated)
		_tcscat(tcDSN,_T(";Trusted_Connection=yes"));
	else
	{
		_tcscat(tcDSN,_T(";Uid="));
		_tcscat(tcDSN,tcUserName);
		_tcscat(tcDSN,_T(";Pwd="));
		_tcscat(tcDSN,tcPassWord);
	}

	// Allocate an ODBC environment handle
	iReturn = SQLAllocHandle(SQL_HANDLE_ENV, NULL, &henv);
	if(SQL_SUCCESS != iReturn)
	{
		// Failed to allocate environment handle
		Sql_SendErr(srvproc,_T("Failed to allocate ODBC environment handle."));
		goto CleanUp;
	}
	SQLSetEnvAttr(henv,SQL_ATTR_ODBC_VERSION,(SQLPOINTER) SQL_OV_ODBC3,SQL_IS_INTEGER);

	// Allocate an ODBC connection handle
	iReturn = SQLAllocHandle(SQL_HANDLE_DBC, henv, &hdbc);
    if(SQL_SUCCESS != iReturn)
	{
		// Failed to allocate connection handle
		Sql_SendErr(srvproc,_T("Failed to allocate ODBC connection handle."));
		goto CleanUp;
	}

	// Set integrated security option if needed
	if(bImpersonated)
		 SQLSetConnectAttr(hdbc,SQL_INTEGRATED_SECURITY,(SQLPOINTER)SQL_IS_ON,SQL_IS_INTEGER);

	// Connect
	iReturn = SQLDriverConnect(
				   hdbc,
				   NULL,
				   (SQLTCHAR*)tcDSN,
				   (SQLSMALLINT)_tcsclen((const TCHAR*)tcDSN),
				   NULL,
				   NULL, 
				   &iOutConnStrLen,
				   SQL_DRIVER_NOPROMPT);

	if(SQL_ERROR == iReturn)
	{
		// Failed to connect
		Sql_SendErr(srvproc,_T("SQLConnect() failed."));
		goto CleanUp;
	}
	
	// Create a statement handle
	iReturn = SQLAllocHandle(SQL_HANDLE_STMT, hdbc, &hstmt);
    if (SQL_SUCCESS != iReturn)
	{
        // Failed to create a statement handle
		Sql_SendErr(srvproc,_T("Failed to create statement handle."));
		goto CleanUp;
	}

	// - Get DLL data from SQL -

	// Get the table name from the xsproc parameters
	if(!GetParamInfo(srvproc,1,&tcTableName))
	{
		// Failed to get the table name parameter
		Sql_SendErr(srvproc,_T("Failed to get the table name parameter information."));
		goto CleanUp;
	}

	// Get the library name from the xsproc parameters
	if(!GetParamInfo(srvproc,2,&tcLibrary))
	{
		// Failed to get the library name parameter
		Sql_SendErr(srvproc,_T("Failed to get the library name parameter information."));
		goto CleanUp;
	}

	// Get the function name from the xsproc parameters
	if(!GetParamInfo(srvproc,3,&tcFunction))
	{
		// Failed to get the function name parameter
		Sql_SendErr(srvproc,_T("Failed to get the function name parameter information."));
		goto CleanUp;
	}

	// ..finally, get the any information for DLL parameters
	if(4 >= srv_rpcparams(srvproc))
	{
		if(!GetParamInfo(srvproc,4,&tcParameters))
		{
			// Failed to get DLL parameters
			Sql_SendErr(srvproc,_T("Failed to get DLL parameter information."));
			goto CleanUp;
		}
	}

	// Build the SQL statement to get the DLL binary data out of the table
	tcSqlStmnt = new TCHAR[_tcslen(_T("SELECT data FROM ")) +
							_tcslen(tcTableName) + 1 +
							_tcslen(_T("WHERE name=?")) + 1];
	memset(tcSqlStmnt,0,(_tcslen(_T("SELECT data FROM ")) + _tcslen(tcTableName) + _tcslen(_T("WHERE name=?")) + 2) * sizeof(TCHAR));
	_tcscpy(tcSqlStmnt,_T("SELECT data FROM "));	//
	_tcscat(tcSqlStmnt,tcTableName);				// Hey look, SQL injection
	_tcscat(tcSqlStmnt,_T(" WHERE name=?"));		//

	// Create the prepared statement
	iReturn = SQLPrepare(hstmt,(SQLTCHAR*)tcSqlStmnt,_tcslen(tcSqlStmnt));
	if(SQL_SUCCESS != iReturn)
	{
		// Failed to prepare statement
		Sql_SendErr(srvproc,_T("Failed to prepare select statement."));
		goto CleanUp;
	}

	// Bind the input statement parameter
	iReturn = SQLBindParameter(hstmt, 1, SQL_PARAM_INPUT, SQL_C_TCHAR,SQL_VARCHAR, 255, 0, tcLibrary, 256, NULL);
	if(SQL_SUCCESS != iReturn)
	{
		// Failed to bind the parameter
		Sql_SendErr(srvproc,_T("Failed to bind select statement parameter."));
		goto CleanUp;
	}

	// Execute the statement
	iReturn = SQLExecute(hstmt);
	if(SQL_ERROR == iReturn)
	{
		// Failed to get the data
		Sql_SendErr(srvproc,_T("Failed to get the DLL data."));
		goto CleanUp;
	}

	// Set initial buffer size for the DLL bytes
	bpDLLData = new BYTE[1024];
	memset(bpDLLData,0,1024);

	// Set receive buffer
	bpSqlData = new BYTE[1024];
	memset(bpSqlData,0,1024);

	// Only concerned about the first row
	iReturn = SQLFetch(hstmt);
	if(SQL_NO_DATA == iReturn)
	{
		// Nothing to fetch
		Sql_SendErr(srvproc,_T("No DLL data was retrieved."));
		goto CleanUp;
	}

	// Get the DLL data in multi-part chunks, resizing the buffer as needed
	while(SQL_NO_DATA != iReturn)
	{
		iReturn = SQLGetData(hstmt, 1, SQL_C_BINARY, bpSqlData, 1024,&iBinaryLenOrInd);

		if((iBinaryLenOrInd > 1024) || (iBinaryLenOrInd == SQL_NO_TOTAL))
			iNumBytes = 1024;
		else
			iNumBytes = iBinaryLenOrInd;

		if(0 == dwDLLSize)
		{
			dwDLLSize = (DWORD)iNumBytes;// + 1;
			bpDLLData = new BYTE[dwDLLSize];
			memset(bpDLLData,0,dwDLLSize);
			memcpy(bpDLLData,bpSqlData,dwDLLSize);// - 1);
		}
		else
		{
			bpDLLTemp = new BYTE[dwDLLSize];
			memset(bpDLLTemp,0,dwDLLSize);
			memcpy(bpDLLTemp,bpDLLData,dwDLLSize);
			delete [] bpDLLData;
			bpDLLData = new BYTE[dwDLLSize + (DWORD)iNumBytes];// + 1];
			memset(bpDLLData,0,(dwDLLSize + (DWORD)iNumBytes));// + 1);
			memcpy(bpDLLData,bpDLLTemp,dwDLLSize);
			delete [] bpDLLTemp;
			memcpy(&bpDLLData[dwDLLSize],bpSqlData,iNumBytes);
			dwDLLSize += (DWORD)iNumBytes;// + 1;
		}
		memset(bpSqlData,0,iNumBytes);
	}
	delete [] bpSqlData;
	SQLFreeStmt(hstmt,SQL_RESET_PARAMS);
	delete [] tcSqlStmnt;
	delete [] tcTableName;
	SQLFreeHandle(SQL_HANDLE_STMT,hstmt);
	SQLDisconnect(hdbc);
	SQLFreeHandle(SQL_HANDLE_DBC, hdbc);
	SQLFreeEnv(henv);

	// If we're still impersonating, end impersonation
	if(bImpersonated)
		srv_revert_to_self(srvproc);

	// - Load the DLL into memory for execution (finally!)

	module = MemoryLoadLibrary((void*)bpDLLData);
	if(NULL == module)
	{
		// Failed to load
		Sql_SendErr(srvproc,_T("Failed to load the DLL module."));
		delete [] bpDLLData;
		return(FAIL);
	}

	// The DLL loader used here requires CHAR rather than WCHAR, perform
	// conversion if needed.
	acFunction = new CHAR[_tcslen(tcFunction) + 1];
	if(2 == sizeof(TCHAR))
	{
		memset(acFunction,0,(_tcslen(tcFunction) + 1));
		wcstombs_s((size_t*)&dwCharsConverted,(char*)acFunction,wcslen((wchar_t*)tcFunction) + 1,(wchar_t*)tcFunction,wcslen((wchar_t*)tcFunction));
	}
	else
		strcpy(acFunction,(CHAR*)tcFunction);

	// Specify the function pointer and get the address
	DLLFromMemory DLLModuleFunction;
	DLLModuleFunction = (DLLFromMemory)MemoryGetProcAddress(module, acFunction);
	if(NULL == DLLModuleFunction)
	{
		Sql_SendErr(srvproc,_T("Couldn't locate module function."));
		return(XP_ERROR);
	}
	
	// Execute the module with parameters if they were provided
	if(!DLLModuleFunction(srvproc, tcParameters))
	{
		Sql_SendErr(srvproc,_T("The library returned a failure."));
		return(XP_ERROR);
	}
	
	// Unload the module
	MemoryFreeLibrary(module);

	// Minor clean-up
	delete [] bpDLLData;
	delete [] tcLibrary;
	delete [] tcFunction;
	if(NULL != acFunction)
			delete [] acFunction;

	return(XP_NOERROR);

CleanUp:
	if(SQL_NULL_HSTMT != hstmt)
	{
		SQLFreeHandle(SQL_HANDLE_STMT,hstmt);
	}
	if(SQL_NULL_HDBC != hdbc)
	{
		SQLDisconnect(hdbc);
		SQLFreeHandle(SQL_HANDLE_DBC, hdbc);
	}

	if(SQL_NULL_HENV != henv)
		SQLFreeEnv(henv);
	return(XP_ERROR);
}

// Prints the xsproc help information
void PrintHelpInfo(SRV_PROC *srvproc)
{
	Sql_SendErr(srvproc,_T("This xsproc requires 3 parameters: [Table Name] [DLL Name] [Function Name] <Function Parameters>"));
	Sql_SendErr(srvproc,_T("-"));
	Sql_SendErr(srvproc,_T("If the table resides in a your non-default database, specify the database.table"));
		return;
}

// Gets textual data from the srv_pfield function
bool GetPfieldInfo(SRV_PROC *srvproc, int iField, TCHAR** tcBuffer)
{
	int		iSize = 0;
	DWORD	dwCharsConverted = 0;

	// Get request pfield info
	srv_pfield(srvproc,iField,&iSize);
	if(0 < iSize)
	{
		*tcBuffer = new TCHAR[iSize + 1];
		memset(*tcBuffer,0,(iSize + 1) * sizeof(TCHAR));

		if(2 == sizeof(TCHAR))
			mbstowcs_s((size_t*)&dwCharsConverted,(wchar_t*)*tcBuffer,(size_t)(iSize + 1),srv_pfield(srvproc,iField,&iSize),(size_t)(iSize));
		else
			strcpy_s((char*)*tcBuffer,(rsize_t)iSize + 1,srv_pfield(srvproc,iField,&iSize));
	}
	
	return(true);
}

// Gets parameter information using srv_paraminfo function
bool GetParamInfo(SRV_PROC *srvproc, int iParam, TCHAR** tcBuffer)
{
	BYTE*	pbData			= NULL;
	BYTE	pbType			= 0;
	ULONG	pcbMaxLen		= 0;
	ULONG	pcbActualLen	= 0;
	BOOL	pfNull			= false;

	if(SUCCEED != srv_paraminfo(srvproc,iParam,&pbType,&pcbMaxLen,&pcbActualLen,NULL,&pfNull))
		return(false);

	// Allocate space for the parameter
	pbData = (BYTE*)malloc(pcbActualLen + 2);
	memset(pbData,0,(pcbActualLen + 2));
	pcbMaxLen = pcbActualLen + 2;

	// Copy the parameter
	if(SUCCEED != srv_paraminfo(srvproc,iParam,&pbType,&pcbMaxLen,&pcbActualLen,pbData,&pfNull))
		return(false);

	if(!NormalizeSqlData(pbData,tcBuffer,pcbActualLen,pbType))
		return(false);

	free(pbData);

	return(true);
}

// Determines the data returned from SQL and puts it in the current working TCHAR format
bool NormalizeSqlData(void* pbData, TCHAR** tczClientRequest, ULONG pcbActualLen, BYTE pbType)
{
	DWORD	dwReturnVal;

	// Determine the data type provided
	switch(pbType)
	{
	// ASCII types
	case SRVBIGCHAR:
	case SRVBIGVARCHAR:
	case SRVCHAR:
	case SRVTEXT:
	case SRVVARCHAR:
		*tczClientRequest = (TCHAR*)malloc((sizeof(TCHAR) * (strlen((char*)pbData) + 1)));
		memset(*tczClientRequest,0,(sizeof(TCHAR) * (strlen((char*)pbData) + 1)));
		if(2 == sizeof(TCHAR))
		{
			// UNICODE to ASCII
			mbstowcs_s((size_t*)&dwReturnVal,(wchar_t*)*tczClientRequest,strlen((char*)pbData) + 1,(char*)pbData,strlen((char*)pbData));
		}
		else
		{
			// ASCII to ASCII
			strcpy_s((char*)*tczClientRequest,strlen((char*)pbData) + 1,(char*)pbData);
		}
		break;
	// UNICODE types
	case SRVNCHAR:
	case SRVNTEXT:
	case SRVNVARCHAR:
		*tczClientRequest = (TCHAR*)malloc((sizeof(TCHAR) * (wcslen((wchar_t*)pbData) + 1)));
		memset(*tczClientRequest,0,(sizeof(TCHAR) * (wcslen((wchar_t*)pbData) + 1)));
		if(2 == sizeof(TCHAR))
		{
			// UNICODE to UNICODE
			wcscpy_s((wchar_t*)*tczClientRequest,wcslen((wchar_t*)pbData) + 1,(wchar_t*)pbData);
		}
		else
		{
			// UNICODE to ASCII
			wcstombs_s((size_t*)&dwReturnVal,(char*)*tczClientRequest,wcslen((wchar_t*)pbData) + 1,(wchar_t*)pbData,wcslen((wchar_t*)pbData));
		}
		break;
	// Non-textual types
	default:
		return(false);
		break;
	}
	return(true);
}

